!! Copyright (C) 2009,2010,2011,2012  Marco Restelli
!!
!! This file is part of:
!!   FEMilaro -- Finite Element Method toolkit
!!
!! FEMilaro is free software; you can redistribute it and/or modify it
!! under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 3 of the License, or
!! (at your option) any later version.
!!
!! FEMilaro is distributed in the hope that it will be useful, but
!! WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
!! General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with FEMilaro; If not, see <http://www.gnu.org/licenses/>.
!!
!! author: Marco Restelli                   <marco.restelli@gmail.com>

!>\brief
!! 
!! Viscous flux, prescribed scalar viscosity
!!
!! \n
!!
!! The energy and momentum diffusive fluxes are
!! \f{displaymath}{
!!  \underline{\underline{\mathcal{F}}}^d_{\underline{U}} =
!!  -\rho\nu\left[ \nabla\underline{u} + \nabla\underline{u}^T +
!!  \lambda\nabla\cdot\underline{u}\,\mathcal{I}\right], \qquad
!!  \underline{\mathcal{F}}^d_E = -\rho\frac{\nu c_p}{Pr}\nabla T +
!!  \underline{u}\cdot
!!  \underline{\underline{\mathcal{F}}}^d_{\underline{U}}.
!! \f}
!<----------------------------------------------------------------------
module mod_viscous_flux

!-----------------------------------------------------------------------

 use mod_kinds, only: &
   mod_kinds_initialized, &
   wp

 use mod_messages, only: &
   mod_messages_initialized, &
   error, warning, info

 use mod_grid, only: &
   mod_grid_initialized, &
   t_grid

 use mod_base, only: &
   mod_base_initialized, &
   t_base

 use mod_atm_refstate, only: &
   mod_atm_refstate_initialized, &
   t_atm_refstate
 
 use mod_turb_flux_base, only: &
   mod_turb_flux_base_initialized, &
   c_turbmod, c_turbmod_progs, c_turbmod_diags, &
   t_turb_diags

 use mod_dgcomp_testcases, only: &
   mod_dgcomp_testcases_initialized, &
   t_phc, phc, coeff_visc

!-----------------------------------------------------------------------
 
 implicit none

!-----------------------------------------------------------------------

! Module interface

 public :: &
   mod_viscous_flux_constructor, &
   mod_viscous_flux_destructor,  &
   mod_viscous_flux_initialized, &
   t_viscous_flux, vf_flux, compute_vf_flux

 private

!-----------------------------------------------------------------------

! Module types and parameters

 ! public members

 !> Viscous flux
 !!
 !! The field \c grad stores the gradients of the following quatities:
 !! <ul>
 !!  <li> temperature (second index equal to 1)
 !!  <li> velocity    (second index from 2 to <tt>1+grid\%d</tt>)
 !!  <li> tracers     (second index \f$\geq\f$<tt>2+grid\%d</tt>)
 !! </ul>
 !! Some private fields are also included for the computation of the
 !! diagnostics.
 type, extends(c_turbmod) :: t_viscous_flux
  real(wp), allocatable :: wg(:) !< Gauss weights, copied from base
  real(wp), allocatable :: base_p(:,:) !< basis functions
 contains 
  procedure, pass(tm) :: init                => vf_init
  procedure, pass(tm) :: clean               => vf_clean
  procedure, pass(tm) :: compute_grad_diags  => vf_compute_grad_diags
  procedure, pass(tm) :: compute_coeff_diags => vf_compute_coeff_diags
  procedure, pass(tm) :: flux                => vf_flux
 end type t_viscous_flux

 ! No need to extend c_turbmod_progs and c_turbmod_diags

! Module variables
 ! public members
 logical, protected ::               &
   mod_viscous_flux_initialized = .false.
 character(len=*), parameter :: &
   this_mod_name = 'mod_viscous_flux'

!-----------------------------------------------------------------------

contains

!-----------------------------------------------------------------------

 subroutine mod_viscous_flux_constructor()
  character(len=*), parameter :: &
    this_sub_name = 'constructor'

   !Consistency checks ---------------------------
   if( (mod_kinds_initialized.eqv..false.) .or. &
    (mod_messages_initialized.eqv..false.) .or. &
        (mod_grid_initialized.eqv..false.) .or. &
        (mod_base_initialized.eqv..false.) .or. &
(mod_atm_refstate_initialized.eqv..false.) .or. &
(mod_turb_flux_base_initialized.eqv..false.).or.&
(mod_dgcomp_testcases_initialized.eqv..false.) ) then
     call error(this_sub_name,this_mod_name, &
                'Not all the required modules are initialized.')
   endif
   if(mod_viscous_flux_initialized.eqv..true.) then
     call warning(this_sub_name,this_mod_name, &
                  'Module is already initialized.')
   endif
   !----------------------------------------------

   mod_viscous_flux_initialized = .true.
 end subroutine mod_viscous_flux_constructor

!-----------------------------------------------------------------------
 
 subroutine mod_viscous_flux_destructor()
  character(len=*), parameter :: &
    this_sub_name = 'destructor'
   
   !Consistency checks ---------------------------
   if(mod_viscous_flux_initialized.eqv..false.) then
     call error(this_sub_name,this_mod_name, &
                'This module is not initialized.')
   endif
   !----------------------------------------------

   mod_viscous_flux_initialized = .false.
 end subroutine mod_viscous_flux_destructor

!-----------------------------------------------------------------------

 pure subroutine vf_init(tm,progs,diags,td,grid,base,ntrcs)
  type(t_grid),       intent(in) :: grid
  type(t_base),       intent(in) :: base
  integer,            intent(in) :: ntrcs
  class(t_viscous_flux), intent(inout) :: tm
  class(c_turbmod_progs), allocatable, intent(out) :: progs
  class(c_turbmod_diags), allocatable, intent(out) :: diags
  type(t_turb_diags), intent(out) :: td

   ! model
   tm%d           = grid%d
   tm%ntrcs       = ntrcs
   tm%initialized = .true.
   allocate(tm%wg(base%m));             tm%wg     = base%wg
   allocate(tm%base_p(base%pk,base%m)); tm%base_p = base%p

   ! no need to allocate progs

   ! diags: the basic type suffices
   allocate(c_turbmod_diags::diags)
   diags%ngrad    = 1 + tm%d + tm%ntrcs
   allocate(diags%grad(tm%d,diags%ngrad,base%pk,grid%ne))

   ! diagnostics
   td%ndiags = 2
   allocate( td%diag_names(td%ndiags) ,  &
     td%diags(td%ndiags,base%pk,grid%ne) )
   td%diag_names(1) = '|nabla^S u + lambda*div(u)|'
   td%diag_names(2) = 'viscous dissipation'

 end subroutine vf_init

!-----------------------------------------------------------------------

 pure subroutine vf_clean(tm,diags,td)
  class(t_viscous_flux), intent(inout) :: tm
  class(c_turbmod_diags), allocatable, intent(inout) :: diags
  type(t_turb_diags), intent(out) :: td
   
   deallocate(tm%wg,tm%base_p)
   deallocate(diags) ! deallocates all the components
   ! td components deallocated automatically

 end subroutine vf_clean

!-----------------------------------------------------------------------

 pure subroutine vf_compute_grad_diags(tm,gd , consv,atm_ref)
  class(t_viscous_flux), intent(in)  :: tm
  class(t_atm_refstate), intent(in)  :: atm_ref(:)
  real(wp),              intent(in)  :: consv(:,:)
  real(wp),              intent(out) :: gd(:,:)

  integer :: id, it
  real(wp), dimension(size(consv,2)) :: rho, e, t
  real(wp), dimension(tm%d,size(consv,2)) :: u, uu

   ! include the reference state
   rho = consv(1,:) + atm_ref%rho
   e   = consv(2,:) + atm_ref%e
   u   = consv(2+1:2+tm%d,:)

   ! velocity
   do id=1,tm%d
     uu(id,:) = u(id,:)/rho
   enddo
   ! temperature
   t = (e/rho - 0.5_wp*sum(uu**2,1) - atm_ref%phi) / phc%cv

   ! collect values
   gd(   1      ,:) = t
   gd(1+1:1+tm%d,:) = uu
   do it=1,tm%ntrcs ! tracers
     gd(1+tm%d+it,:) = consv(1+tm%d+it,:)/rho
   enddo
 
 end subroutine vf_compute_grad_diags

!-----------------------------------------------------------------------

 pure subroutine vf_compute_coeff_diags(tm,cd , grid,base, &
                                        uuu,progs,atm_ref)
  class(t_viscous_flux),  intent(in)    :: tm
  type(t_grid),           intent(in)    :: grid
  type(t_base),           intent(in)    :: base   
  real(wp),               intent(in)    :: uuu(:,:,:)
  class(t_atm_refstate),  intent(in)    :: atm_ref(:,:)
  class(c_turbmod_progs), allocatable, intent(in) :: progs
  class(c_turbmod_diags), intent(inout) :: cd

   !> Nothing to do

 end subroutine vf_compute_coeff_diags

!-----------------------------------------------------------------------

 !> Compute the viscous flux
 !!
 !! This subroutine is split into two to expose a more detailed
 !! interface which can be used to recover the intermediate quantities
 !! computed here. The detailed interface can be useful in the
 !! implementation of other dissipative fluxes which are based on the
 !! viscous one, such as the Smagorinsky flux.
 pure subroutine vf_flux(fem, tm,ie,x,bp, consv,atm_ref, rho,p,uu,cc, &
                         progs,diags, td)
  class(t_viscous_flux), intent(in) :: tm
  integer,               intent(in) :: ie
  real(wp),              intent(in) :: x(:,:)  !< coordinates
  real(wp),              intent(in) :: bp(:,:) !< basis functions
  real(wp),              intent(in) :: consv(:,:)
  class(t_atm_refstate), intent(in) :: atm_ref(:)
  real(wp),              intent(in) :: rho(:), p(:), uu(:,:), cc(:,:)
  class(c_turbmod_progs), allocatable, intent(in) :: progs
  class(c_turbmod_diags), intent(in) :: diags
  real(wp),              intent(out) :: fem(:,:,:)
  type(t_turb_diags),    intent(inout), optional :: td

   call compute_vf_flux(fem, tm,ie,x,bp, consv,atm_ref, &
                        rho,p,uu,cc, diags, td=td)

 end subroutine vf_flux

 pure subroutine compute_vf_flux(fem, tm,ie,x,bp, consv,atm_ref, &
                          rho,p,uu,cc, diags, td, vf_sym,vf_gradg)
  class(t_viscous_flux), intent(in) :: tm
  integer,               intent(in) :: ie
  real(wp),              intent(in) :: x(:,:)
  real(wp),              intent(in) :: bp(:,:)
  real(wp),              intent(in) :: consv(:,:)
  class(t_atm_refstate), intent(in) :: atm_ref(:)
  real(wp),              intent(in) :: rho(:), p(:), uu(:,:), cc(:,:)
  class(c_turbmod_diags), intent(in) :: diags
  real(wp),              intent(out) :: fem(:,:,:)
  type(t_turb_diags),    intent(inout), optional :: td
  real(wp), intent(out), optional :: vf_sym(:,:,:), vf_gradg(:,:,:)

  integer :: i, l, id, jd
  real(wp) :: &
   nu(2,size(x,2)), & !< viscosity
   gradg(tm%d,diags%ngrad,size(x,2)), & !< gradient in physical space
   gsuu (tm%d,   tm%d    ,size(x,2)), & !< velocity symmetric gradient
   ldivu, sym(tm%d,tm%d), &
   s_abs(size(x,2)), diss(size(x,2))
 
   ! problem coefficients
   nu = coeff_visc( x , rho,p,p/(rho*phc%rgas) )

   ! gradients
   !do i=1,diags%ngrad
   !  gradg(:,i,:) = matmul(diags%grad(:,i,:,ie),bp)
   !enddo
   gradg = 0.0_wp ! equivalent implementation
   do l=1,size(bp,2)
     do i=1,size(bp,1)
       gradg(:,:,l) = gradg(:,:,l) + diags%grad(:,:,i,ie)*bp(i,l)
     enddo
   enddo
   if(present(vf_gradg)) vf_gradg = gradg
   do id=1,tm%d
     do jd=1,id-1; gsuu(id,jd,:) = gsuu(jd,id,:); enddo ! symmetry
     do jd=id,tm%d
       gsuu(id,jd,:) = gradg(jd,1+id,:)+gradg(id,1+jd,:)
     enddo
   enddo

   do l=1,size(x,2)
     ldivu = (-2.0_wp/3.0_wp)*0.5_wp*tr(gsuu(:,:,l)) ! lambda*div(u)
     sym = gsuu(:,:,l) ! symmetric gradient
     do id=1,tm%d
       sym(id,id) = sym(id,id) + ldivu
     enddo
     if(present(vf_sym)) vf_sym(:,:,l) = sym

     ! momentum flux
     fem(:,2:1+tm%d,l) = -rho(l)*nu(2,l)*sym

     ! energy flux
     fem(:,   1    ,l) = -rho(l)*nu(1,l)*phc%cp*gradg(:,1,l)  &
                        + matmul( uu(:,l) , fem(:,2:1+tm%d,l) )

     ! tracers
     fem(:,2+tm%d: ,l) = -rho(l)*nu(1,l)*gradg(:,2+tm%d:,l)
     
     if(present(td)) then
       s_abs(l) = sqrt(sum(sym**2))
       diss (l) = -sum(fem(:,2:1+tm%d,l)*sym) ! use sym^T=sym
     endif
   enddo

   ! the projections are simple because the basis is orthonormal
   if(present(td)) then
     td%diags(1,:,ie) = matmul( tm%base_p , tm%wg*s_abs )
     td%diags(2,:,ie) = matmul( tm%base_p , tm%wg*diss  )
   endif

 end subroutine compute_vf_flux
 
!-----------------------------------------------------------------------

 pure function tr(a)
  real(wp), intent(in) :: a(:,:)
  real(wp) :: tr
  integer :: i

   tr = a(1,1)
   do i=2,minval(shape(a))
     tr = tr + a(i,i)
   enddo
 end function tr

!-----------------------------------------------------------------------

end module mod_viscous_flux

